fix(s3): stream wr.s3.download in chunks instead of loading whole object into memory - #3396
Closed
nileshpatil6 wants to merge 1 commit into
Closed
fix(s3): stream wr.s3.download in chunks instead of loading whole object into memory#3396nileshpatil6 wants to merge 1 commit into
nileshpatil6 wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes/relates to #2831.
wr.s3.download currently opens the object with s3_block_size=-1 (one shot download), which fetches the entire object into an in memory cache inside _S3ObjectBase, then calls s3_f.read() to pull the whole thing out as a second bytes object, then writes it to the local file. For large files this means the object briefly exists twice in memory before it even reaches disk, so downloading a multi GB file can OOM a process that would otherwise have plenty of headroom for a streamed copy.
This changes download() to open the object with a fixed 8 MB block size instead of one shot, and copies it to the destination in a simple read/write loop (_copy_in_chunks) instead of a single read() call. That keeps memory usage roughly proportional to the block size regardless of the source object size. The underlying _S3ObjectBase.read() already supports incremental ranged reads when s3_block_size is a positive number (this is the same mechanism _read_text_core.py already uses for streaming reads), so this reuses existing, tested code paths rather than adding new S3 fetch logic.
Added test_download_file_chunked in tests/unit/test_moto.py, which monkeypatches the block size down to 5 bytes and downloads an object whose size is not an exact multiple of the block size, to exercise the multi chunk path and confirm the bytes are reassembled correctly.
Testing done locally:
I did not have access to real AWS infra so I couldn't run the basic/full test environments described in CONTRIBUTING.md, only the mocked moto based unit tests.